module ModuleSolveEqm

    use Globals
    use ModuleInitialize
    use ModuleSolve_GLO_LOC
    use ModuleAggregates
    use ModuleSaving
    use ModuleSteady
    use ModuleDistributions

    contains
    
function welfare_eqm(param_in_vec,index_ss)

    implicit none
    
    ! Local variable declarations
    real(8), intent(in) :: param_in_vec(3)    ! input parameters
    integer, intent(in) :: index_ss           ! =1 if initial, 2 if LR
    real(8) :: welfare_eqm                    ! equilibrium welfare
    
    ! Show input parameters
    print *, ''
    print *, 'Solve for steady state given tax parameters'
    print *, 'param_in_vec = ', param_in_vec
    print *, 'size(param_in_vec) = ', size(param_in_vec)
    
    ! Set global variables for tax function equal to inputs
    ! omegat set globally in main file
    gamma   = param_in_vec(1)       ! tax progressivity
    mt      = param_in_vec(2)       ! transfer level
    rt      = param_in_vec(3)       ! transfer phase-out

    ! Prepare Newton
    maxit_r = 25                        ! max number of iterations for asset market clearing
    tol_r = 1D-5                        ! tolerance for asset market clearing
    maxitV = 100                        ! maximum number of iterations for VFI
    ww = 0.2D0                          ! updating speed for VFI
    x0(1) = min(max(lbdlb,lambda),lbdub)    ! bounds on lambda
    x0(2) = min(max(0.005D0,r),0.06D0)  ! bounds on r
    tolf = tol_r*0.5D0                  ! tolerance Newton
    maxitqbs = 15; maxitq = 80          ! max iterations (backstep; Newton)
    max_exit = 40                       ! max iterations

     ! Wage implied by interest rate
    wge = (1.0D0-alpha)/(x0(2)+delta)
    wge = alpha*(wge**((1D0-alpha)/alpha))
    
    ! Solve for equilibrium
    call StartVals          ! creates grids for assets and labor
    call InitialVGuess_wa   ! initial guess for value function; assumption: everybody has same states forever and consumes income
    call InitializeMeasure  ! initialization of stationary measure; equal shares across all states
    call ComputeVE          ! continuation value
    call ComputeSteady      ! solve for stationary equilibrium

    ! Collect outputs
    welfare_eqm = -sum(mu*V)     ! Welfare
    
    if (index_ss==1) then ! initial steady-state
        print *, ''
        print *, 'Initial steady state results'
        print *, ''
    else                  ! final steady-state
        print *, ''
        print *, 'Final steady state results'
        print *, ''
    endif
    
    print *, 'Tax system with gamma = ', gamma, ' , mt = ', mt, ' , rt = ', rt
    print *, 'Equilibrium with lambda = ', lambda, ' , r = ', r
    print *, ''
    print *, 'Steady state welfare =', welfare_eqm
    print *, 'Yagg = ', Yagg, ' ,Hagg = ', Hagg, ' ,Lagg = ', Lagg, ' ,Kagg =', Kagg
    print *, 'G/Y and D/Y in % = ', (G/Yagg)*100d0, (Debt/Yagg)*100d0
    print *, '(K+D)/Y in % = ', ((Kagg + Debt)/Yagg)*100d0
    print *, 'Net interest expenditure in % ', ((r*Debt)/Yagg)*100d0
    print *, 'ymean, ymean_new = ', ymean, ymean_new
    print *, ''
    print *, 'Income distribution'
    call income_distribution
    print *, ''
    print *, 'Income growth distribution'
    call income_growth
    ! Call implied rates calculation
    call implied_rates(index_ss,1000)
    ! Call tax rates distribution calculation
    call tax_distribution(index_ss)

    ! Store errors in global variable
    error_vec(1) = abs(asset)   ! Error asset market clearing
    error_vec(2) = abs(BC)      ! Error govt. budget clearing
    error_vec(3) = errV         ! Error VFI
    error_vec(4) = errmu        ! Error measure

    ! Store policy functions
    if (index_ss==1) then ! initial steady-state
        a_pol_init = a_pol
        c_pol_init = c_pol
        h_pol_init = h_pol
        mu_init    = mu
        call SavePolicies_INIT
        call SavePrices_INIT
    else                  ! final steady-state
        a_pol_lr = a_pol
        c_pol_lr = c_pol
        h_pol_lr = h_pol
        mu_lr    = mu
        V_lr     = V
        call SavePolicies_LR
        call SavePrices_LR
    endif

end function

end module
